# -*- coding: utf-8 -*-

# Copyright (C) 2011 Plamen Valov


import gdb

import sys
import re


class MEM_SECT_FLAG:
    (
     ALLOC,     ## Section will have space allocated in the process when loaded.
                ## Set for all sections except those containing debug information
     
     LOAD,      ## Section will be loaded from the file into the child process memory.
                ## Set for pre-initialized code and data, clear for .bss sections
     
     RELOC,     ## Section needs to be relocated before loading
     
     READONLY,  ## Section cannot be modified by the child proces
     
     CODE,      ## Section contains executable code only
     
     DATA,      ## Section contains data only (no executable code)
     
     ROM,       ## Section will reside in ROM
     
     CONSTRUCTOR,       ## Section contains data for constructor/destructor lists
     
     HAS_CONTENTS,      ## Section is not empty
     
     NEVER_LOAD,        ## An instruction to the linker to not output the section
     
     COFF_SHARED_LIBRARY,       ## A notification to the linker that the section
                                ## contains COFF shared library information
     
     IS_COMMON      ## Section contains common symbols
    ) = xrange(12)


class Unsigned(object):
    def __new__(self, value, max_bytes=None):
        assert isinstance(value, int) or isinstance(value, long),\
               "Value must be int: %s, %s" % (value, type(value).__name__)
        if(max_bytes is None):
            if(sys.maxsize > 2**32):
                mask = 0xffffffffffffffff
            else:
                mask = 0xffffffff
        else:
            mask = 2**max_bytes - 1
        
        return mask & value


class GdbEerror(StandardError):
    pass


class MemoryAccessError(GdbEerror):
    pass


class SectionsError(MemoryAccessError):
    pass


class MemSectionSearcher(object):
    def __init__(self, inferior_id=None):
        try:
            self._inferior = gdb.inferiors()[0]
        except Exception, err:
            raise GdbEerror("Could not get inferior: %s" % str(err))
        self._sections = self._parse_sections_info(self._get_sections())
    
    def _get_sections(self):
        sections_data = gdb.execute("maintenance info sections ALLOBJ", to_string=True).strip()
        
        if(not sections_data):
            raise SectionsError("Could not get sections - probably no process loaded")
        
        return sections_data
    
    def _parse_sect_flags(self, flag_data):
        flags = set()
        flag_list = [fl.strip().upper() for fl in flag_data.strip().split(' ')]
        for fl in flag_list:
            if(fl == 'ALLOC'):
                flags.add('ALLOC')
            elif(fl == 'LOAD'):
                flags.add('LOAD')
            elif(fl == 'RELOC'):
                flags.add('RELOC')
            elif(fl == 'READONLY'):
                flags.add('READONLY')
            elif(fl == 'CODE'):
                flags.add('CODE')
            elif(fl == 'DATA'):
                flags.add('DATA')
            elif(fl == 'ROM'):
                flags.add('ROM')
            elif(fl == 'CONSTRUCTOR'):
                flags.add('CONSTRUCTOR')
            elif(fl == 'HAS_CONTENTS'):
                flags.add('HAS_CONTENTS')
            elif(fl == 'NEVER_LOAD'):
                flags.add('NEVER_LOAD')
            elif(fl == 'COFF_SHARED_LIBRARY'):
                flags.add('COFF_SHARED_LIBRARY')
            elif(fl == 'IS_COMMON'):
                flags.add('IS_COMMON')
            else:
                raise SectionsError("Could not parse flags: %s" % flag_data)
        
        return flags
    
    def _parse_sections_info(self, data):
        sections = []
        
        objects = [dt.strip() for dt in data.split("Object file:") if dt.strip()]
        objects.pop(0)
        
        for section_dt in objects:
            sect = {}
            
            scts = [i.strip() for i in section_dt.split("\n")]
            obj_file = scts.pop(0)
            
            sect['objfile'] = obj_file
            sect['sections'] = []
            sections.append(sect)
            
            ## Example line:
            ## 0x80482a8->0x80482b0 at 0x000002a8: .rel.dyn ALLOC LOAD READONLY DATA HAS_CONTENTS
            sect_pattern = "^" + "(0x[A-Fa-f0-9]+)"\
                           + re.escape("->")\
                           + "(0x[A-Fa-f0-9]+)"\
                           + " .*?"\
                           + re.escape(": ")\
                           + "([^\s]+)" + "(.*?)" + "$"
            sect_regex = re.compile(sect_pattern)
            
            for sct in scts:
                match = sect_regex.match(sct)
                if(match is None):
                    raise SectionsError("Could not parse section info: %s" % sct)
                
                start_addr = Unsigned(int(match.group(1), base=16))
                end_addr = Unsigned(int(match.group(2), base=16))
                sect_name = match.group(3).strip()
                sect_flags = self._parse_sect_flags(match.group(4).strip())
                
                sect['sections'].append({'name': sect_name,
                                         'start': start_addr,
                                         'end': end_addr,
                                         'flags': sect_flags})
        
        return sections
    
    def searchiter(self, pattern, section_flags=frozenset(), objfile=None, section=None):
        matches = []
        pattern_len = len(pattern)
        for obj in self._sections:
            obj_name = obj['objfile']
            for sect in obj['sections']:
                start_addr = sect['start']
                while(True):
                    match_val = self._inferior.search_memory(start_addr, sect['end'] - start_addr, pattern)
                    if(match_val is None):
                        break
                    else:
                        match_addr = Unsigned(match_val)
                        yield (obj_name, sect['name'], sect['flags'], match_addr)
                        start_addr = match_addr + pattern_len


def search(pattern, section_flags=frozenset()):
    searcher = MemSectionSearcher()
    print
    for match in searcher.searchiter(pattern, section_flags):
        print "%s --> %s [%s]: 0x%x" % (match[0], match[1], ' '.join([i for i in match[2]]), match[3])
    print
